use comp_agent::controller::ssh_execute_shell::{SshExecutor, ExecutionStatus, ExecutionResult, EXECUTIONS};
use std::{sync::{Arc, Mutex}, time::Duration};
use tokio;
use uuid::Uuid;

// 重置全局执行结果存储
fn reset_executions() {
    let mut executions = EXECUTIONS.lock().unwrap();
    executions.clear();
}

#[cfg(test)]
mod unit_test {
    use super::*;

    // 测试常量
    const TEST_HOST: &str = "10.86.13.11";
    const TEST_PORT: u16 = 31222;
    const TEST_USER: &str = "app";
    const TEST_PASSWORD: &str = "app@123";
    const TEST_TARGET_USER: &str = "root";
    const TIMEOUT: u64 = 10;
    const MAX_OUTPUT: usize = 3 * 1024 * 1024;

    // 备用测试配置
    // const TEST_HOST: &str = "10.88.152.30";
    // const TEST_PORT: u16 = 3022;
    // const TEST_USER: &str = "o2ma";
    // const TEST_PASSWORD: &str = "Mi@Np4Tf!Hd6";
    // const TEST_TARGET_USER: &str = "o2ma";
    // const TIMEOUT: u64 = 10;

    fn create_test_executor() -> Result<SshExecutor, Box<dyn std::error::Error>> {
        println!("Connecting to {}:{} as user {}", TEST_HOST, TEST_PORT, TEST_USER);
        let executor = SshExecutor::connect(TEST_HOST, TEST_PORT, TEST_USER, TEST_PASSWORD, TIMEOUT)?;
        println!("Successfully connected to {}:{} as user {}", TEST_HOST, TEST_PORT, TEST_USER);
        Ok(executor)
    }

    // 辅助函数：等待并获取执行结果
    async fn wait_for_execution(executor: &mut SshExecutor, execution_id: &str, wait_secs: u64) -> Result<ExecutionResult, Box<dyn std::error::Error>> {
        println!("Waiting {} seconds for execution {}", wait_secs, execution_id);
        tokio::time::sleep(Duration::from_secs(wait_secs)).await;
        let result = executor.get_execution_result(execution_id).await;
        match &result {
            Ok(r) => {
                println!("Execution {} result: status={:?}, exit_code={:?}, stdout='{}', stderr='{}'",
                    execution_id, r.status, r.exit_code, r.stdout, r.stderr);
            },
            Err(e) => {
                println!("Failed to get result for execution {}: {}", execution_id, e);
            }
        }
        result
    }

    #[tokio::test]
    async fn test_execute() {
        reset_executions();
        
        // 创建SSH执行器
        let mut executor = create_test_executor().expect("Failed to create SSH executor");

        // 执行命令
        let execution_id = executor.execute(None, "echo 'Hello World' && exit 1").await
            .expect("Failed to execute command");

        println!("Execution ID: {}", execution_id);

        // 等待命令执行完成并获取结果
        let result = wait_for_execution(&mut executor, &execution_id, 2).await
            .expect("Failed to get execution result");

        println!("Status: {:?}", result.status);
        println!("Output: {}", result.stdout);
        println!("Error: {}", result.stderr);

        // 断言
        assert_eq!(result.status, ExecutionStatus::Completed);
        assert_eq!(result.id, execution_id);
        assert!(result.stdout.contains("Hello World"));
        assert!(result.stderr.is_empty());
        assert_eq!(result.exit_code, Some(0));
    }

    #[tokio::test]
    async fn test_execute_timeout() {
        reset_executions();
        
        // 创建SSH执行器，设置短超时
        let mut executor = SshExecutor::connect(TEST_HOST, TEST_PORT, TEST_USER, TEST_PASSWORD, 2)
            .expect("Failed to create SSH executor");

        // 执行一个会超时的命令
        let execution_id = executor.execute(None, "sleep 5 && echo 'This should not be printed'").await
            .expect("Failed to execute command");

        // 等待超时发生
        let result = wait_for_execution(&mut executor, &execution_id, 3).await
            .expect("Failed to get execution result");

        println!("Timeout test status: {:?}", result.status);
        println!("Timeout test output: {}", result.stdout);

        // 断言命令超时
        assert_eq!(result.status, ExecutionStatus::Timeout);
        assert_eq!(result.id, execution_id);
        assert!(!result.stdout.contains("This should not be printed"));
        assert!(result.stdout.contains("[错误] 命令在"));
        assert_eq!(result.exit_code, Some(-1));
    }

    #[tokio::test]
    async fn test_execute_interrupted() {
        reset_executions();
        
        // 创建SSH执行器
        let mut executor = SshExecutor::connect(TEST_HOST, TEST_PORT, TEST_USER, TEST_PASSWORD, 10)
            .expect("Failed to create SSH executor");

        // 执行命令
        let execution_id = executor.execute(None, "sleep 5 && echo 'This should not be printed'").await
            .expect("Failed to execute command");

        // 等待一小段时间确保命令已经启动
        tokio::time::sleep(Duration::from_millis(500)).await;
        
        // 手动设置中断标志
        executor.set_interrupted(true);
        
        // 等待命令执行完成
        let result = wait_for_execution(&mut executor, &execution_id, 2).await
            .expect("Failed to get execution result");

        println!("Interrupted test status: {:?}", result.status);
        println!("Interrupted test output: {}", result.stdout);

        // 断言命令被中断
        assert_eq!(result.status, ExecutionStatus::Interrupted);
        assert_eq!(result.id, execution_id);
        assert!(!result.stdout.contains("This should not be printed"));
        assert!(result.stdout.contains("[中断] 命令执行被中断"));
        assert_eq!(result.exit_code, Some(-5));
    }

    #[tokio::test]
    async fn test_execute_large_output() {
        reset_executions();
        
        // 创建SSH执行器
        let mut executor = create_test_executor().expect("Failed to create SSH executor");

        // 执行生成大输出的命令
        // 使用行数代替字节数，确保兼容性
        let execution_id = executor.execute(None, "yes | head -n 100000").await
            .expect("Failed to execute command");

        // 等待命令执行完成
        let result = wait_for_execution(&mut executor, &execution_id, 5).await
            .expect("Failed to get execution result");

        println!("Large output test status: {:?}", result.status);
        println!("Large output size: {} bytes", result.stdout.len());

        // 断言命令成功执行且输出被限制
        assert_eq!(result.status, ExecutionStatus::Completed);
        assert_eq!(result.id, execution_id);
        assert!(result.stdout.len() <= MAX_OUTPUT);
        assert!(result.stdout.contains("[输出限制]"));
    }

    #[tokio::test]
    async fn test_execute_invalid_command() {
        reset_executions();
        
        let mut executor = create_test_executor().expect("Failed to create SSH executor");
        let execution_id = executor.execute(None, "invalid_command").await
            .expect("Failed to execute invalid command");

        // 等待命令执行完成并获取结果
        let result = wait_for_execution(&mut executor, &execution_id, 2).await
            .expect("Failed to get execution result");

        println!("Status: {:?}", result.status);
        println!("Exit code: {:?}", result.exit_code);
        println!("Output: {}", result.stdout);
        println!("Error: {}", result.stderr);

        // 断言
        assert_eq!(result.status, ExecutionStatus::Failed);
        assert_eq!(result.id, execution_id);
        assert!(result.stderr.contains("command not found") || !result.stderr.is_empty());
        assert_eq!(result.exit_code, Some(127));
    }

    #[tokio::test]
    async fn test_execute_as_user_invalid_target() {
        reset_executions();
        
        let mut executor = create_test_executor().expect("Failed to create SSH executor");
        let invalid_user = "non_existent_user_123456";

        // 尝试以不存在的用户身份执行命令
        let execution_id = executor.execute_as_user(None, "whoami", invalid_user, Some(TEST_PASSWORD)).await;

        match execution_id {
            Ok(id) => {
                // 等待命令执行完成
                let result = wait_for_execution(&mut executor, &id, 2).await
                    .expect("Failed to get execution result");

                println!("Invalid user test status: {:?}", result.status);
                println!("Invalid user test error: {}", result.stderr);

                // 断言命令失败
                assert_eq!(result.status, ExecutionStatus::Failed);
                assert_eq!(result.id, id);
                assert!(result.stderr.contains("user not found") || result.stderr.contains("无法找到用户") || !result.stderr.is_empty());
            },
            Err(e) => {
                // 如果在执行阶段就失败，也是预期的
                let error_msg = e.to_string();
                println!("Invalid user test failed during execution: {}", error_msg);
                assert!(error_msg.contains("user not found") || error_msg.contains("无法找到用户") || !error_msg.is_empty());
            }
        }
    }

    #[tokio::test]
    async fn test_execute_as_user_invalid_password() {
        reset_executions();
        
        let mut executor = create_test_executor().expect("Failed to create SSH executor");
        let invalid_password = "wrong_password_123";

        // 使用无效密码尝试sudo
        let execution_id = executor.execute_as_user(None, "whoami", TEST_TARGET_USER, Some(invalid_password)).await;

        match execution_id {
            Ok(id) => {
                // 等待命令执行完成
                let result = wait_for_execution(&mut executor, &id, 2).await
                    .expect("Failed to get execution result");

                println!("Invalid password test status: {:?}", result.status);
                println!("Invalid password test error: {}", result.stderr);

                // 断言命令失败
                assert_eq!(result.status, ExecutionStatus::Failed);
                assert_eq!(result.id, id);
                assert!(result.stderr.contains("incorrect password") || result.stderr.contains("密码错误") || !result.stderr.is_empty());
            },
            Err(e) => {
                // 如果在执行阶段就失败，也是预期的
                let error_msg = e.to_string();
                println!("Invalid password test failed during execution: {}", error_msg);
                assert!(error_msg.contains("incorrect password") || error_msg.contains("密码错误") || !error_msg.is_empty());
            }
        }
    }

    #[tokio::test]
    async fn test_execute_as_user_passwordless_sudo() {
        let mut executor = create_test_executor().expect("Failed to create SSH executor");
        let execution_id = executor.execute_as_user(None, "whoami", TEST_TARGET_USER, None).await;

        match execution_id {
            Ok(id) => {
                // 等待命令执行完成并获取结果
                let result = wait_for_execution(&mut executor, &id, 2).await
                    .expect("Failed to get execution result");

                println!("Passwordless sudo result: status={:?}, output={}", result.status, result.stdout);
                
                // 检查命令是否成功执行
                if result.status == ExecutionStatus::Completed {
                    assert_eq!(result.stdout.trim(), TEST_TARGET_USER);
                } else if result.stdout.contains("sudo: not found") || result.stderr.contains("sudo: not found") {
                    println!("Skipping assertion because sudo is not available on the target server");
                } else {
                    panic!("Command failed with unexpected error: stdout={}, stderr={}", result.stdout, result.stderr);
                }
            }
            Err(e) => {
                // 无密码sudo可能在某些环境中失败，这是预期的
                let error_msg = e.to_string();
                println!("Passwordless sudo failed (expected for some cases): {}", error_msg);
                
                // 检查是否是因为sudo命令不存在导致的失败
                if error_msg.contains("sudo: not found") {
                    println!("Skipping assertion because sudo is not available on the target server");
                } else {
                    panic!("Unexpected failure: {}", e);
                }
            }
        }
    }

    #[tokio::test]
    async fn test_execute_as_user_with_custom_id() {
        let mut executor = create_test_executor().expect("Failed to create SSH executor");
        let custom_id = "test_as_user_custom_id_123";

        // 使用自定义ID执行命令
        let execution_id = executor.execute_as_user(Some(custom_id.to_string()), "whoami", TEST_TARGET_USER, Some(TEST_PASSWORD)).await;

        match execution_id {
            Ok(id) => {
                // 验证ID是否正确
                assert_eq!(id, custom_id);

                // 等待命令执行完成并获取结果
                let result = wait_for_execution(&mut executor, &id, 2).await
                    .expect("Failed to get execution result");

                // 检查命令是否成功执行
                if result.status == ExecutionStatus::Completed {
                    assert_eq!(result.stdout.trim(), TEST_TARGET_USER);
                } else if result.stdout.contains("sudo: not found") || result.stderr.contains("sudo: not found") {
                    println!("Skipping assertion because sudo is not available on the target server");
                }
            },
            Err(e) => {
                let error_msg = e.to_string();
                println!("Execute as user with custom ID failed: {}", error_msg);
                
                // 检查是否是因为sudo命令不存在导致的失败
                if error_msg.contains("sudo: not found") {
                    println!("Skipping assertion because sudo is not available on the target server");
                } else {
                    panic!("Unexpected failure: {}", e);
                }
            }
        }
    }

    #[tokio::test]
    async fn test_execute_as_user_special_chars() {
        let mut executor = create_test_executor().expect("Failed to create SSH executor");
        let special_cmd = "echo 'Special chars: !@#$%^&*()'";

        // 执行包含特殊字符的命令
        let execution_id = executor.execute_as_user(None, special_cmd, TEST_TARGET_USER, Some(TEST_PASSWORD)).await;

        match execution_id {
            Ok(id) => {
                // 等待命令执行完成并获取结果
                let result = wait_for_execution(&mut executor, &id, 2).await
                    .expect("Failed to get execution result");

                println!("Special chars test result: status={:?}, output={}", result.status, result.stdout);
                
                // 检查命令是否成功执行
                if result.status == ExecutionStatus::Completed {
                    assert!(result.stdout.contains("Special chars: !@#$%^&*()"));
                } else if result.stdout.contains("sudo: not found") || result.stderr.contains("sudo: not found") {
                    println!("Skipping assertion because sudo is not available on the target server");
                } else {
                    panic!("Command failed with unexpected error: stdout={}, stderr={}", result.stdout, result.stderr);
                }
            }
            Err(e) => {
                let error_msg = e.to_string();
                println!("Special chars test failed: {}", error_msg);
                
                // 检查是否是因为sudo命令不存在导致的失败
                if error_msg.contains("sudo: not found") {
                    println!("Skipping assertion because sudo is not available on the target server");
                } else {
                    panic!("Unexpected failure: {}", e);
                }
            }
        }
    }

    #[tokio::test]
    async fn test_execute_as_user_sudo_with_password() {
        reset_executions();
        let mut executor = create_test_executor().expect("Failed to create SSH executor");
        let execution_id = executor.execute_as_user(None, "whoami", TEST_TARGET_USER, Some(TEST_PASSWORD)).await
            .expect("Failed to execute command as user with password");

        // 等待命令执行完成并获取结果
        let result = wait_for_execution(&mut executor, &execution_id, 2).await
            .expect("Failed to get execution result");

        println!("Sudo with password result: status={:?}, output={}", result.status, result.stdout);

        // 检查命令是否成功执行
        if result.stdout.contains("sudo: not found") || result.stderr.contains("sudo: not found") {
            println!("Skipping assertion because sudo is not available on the target server");
        } else {
            assert_eq!(result.status, ExecutionStatus::Completed);
            assert_eq!(result.stdout.trim(), TEST_TARGET_USER);
            assert_eq!(result.id, execution_id);
            assert!(result.stderr.is_empty());
        }
    }

    #[tokio::test]
    async fn test_execute_as_user_success() {
        reset_executions();
        let mut executor = create_test_executor().expect("Failed to create SSH executor");

        // 执行命令
        let execution_id = executor.execute_as_user(None, "echo 'Hello from target user'", TEST_TARGET_USER, Some(TEST_PASSWORD)).await
            .expect("Failed to execute command as user");

        // 等待命令执行完成并获取结果
        let result = wait_for_execution(&mut executor, &execution_id, 2).await
            .expect("Failed to get execution result");

        // 断言
        if result.stdout.contains("sudo: not found") || result.stderr.contains("sudo: not found") {
            println!("Skipping assertion because sudo is not available on the target server");
        } else {
            assert_eq!(result.status, ExecutionStatus::Completed);
            assert_eq!(result.id, execution_id);
            assert!(result.stdout.contains("Hello from target user"));
            assert!(result.stderr.is_empty());
        }
    }

    #[tokio::test]
    async fn test_execute_with_custom_id() {
        reset_executions();
        
        let mut executor = create_test_executor().expect("Failed to create SSH executor");
        let custom_id = "test_custom_id_123";
        let execution_id = executor.execute(Some(custom_id.to_string()), "echo 'Custom ID Test'").await
            .expect("Failed to execute command with custom ID");

        assert_eq!(execution_id, custom_id);

        // 等待命令执行完成并获取结果
        let result = wait_for_execution(&mut executor, &execution_id, 2).await
            .expect("Failed to get execution result");

        assert_eq!(result.status, ExecutionStatus::Completed);
        assert_eq!(result.id, custom_id);
        assert!(result.stdout.contains("Custom ID Test"));
    }

    #[tokio::test]
    async fn test_execution_id_generation() {
        reset_executions();
        
        let mut executor = create_test_executor().expect("Failed to create SSH executor");
        
        // 执行两次命令，不提供ID
        let id1 = executor.execute(None, "echo 1").await.expect("Failed to execute command 1");
        let id2 = executor.execute(None, "echo 2").await.expect("Failed to execute command 2");
        
        // 验证生成的ID是有效的UUID v7
        assert!(Uuid::parse_str(&id1).is_ok());
        assert!(Uuid::parse_str(&id2).is_ok());
        
        // 验证两个ID不同
        assert_ne!(id1, id2);
    }

    // 测试报告生成函数
    #[test]
    fn generate_test_report() {
        println!("\n=== SSH Execute Shell Tests Report ===\n");
        println!("Total tests: 12\n");
        println!("1. test_execute: Tests basic command execution functionality");
        println!("   - Verifies: Successful execution, correct output");
        println!("   - Status: PENDING\n");
        println!("2. test_execute_timeout: Tests command timeout functionality");
        println!("   - Verifies: Command times out after specified duration");
        println!("   - Status: PENDING\n");
        println!("3. test_execute_interrupted: Tests command interruption functionality");
        println!("   - Verifies: Command can be interrupted");
        println!("   - Status: PENDING\n");
        println!("4. test_execute_large_output: Tests handling of large output");
        println!("   - Verifies: Output is limited to MAX_OUTPUT");
        println!("   - Status: PENDING\n");
        println!("5. test_execute_invalid_command: Tests execution of invalid command");
        println!("   - Verifies: Proper error handling for invalid commands");
        println!("   - Status: PENDING\n");
        println!("6. test_execute_as_user_passwordless_sudo: Tests passwordless sudo execution");
        println!("   - Verifies: Execution as target user without password");
        println!("   - Status: PENDING\n");
        println!("7. test_execute_as_user_with_custom_id: Tests execute_as_user with custom ID");
        println!("   - Verifies: Custom ID is correctly handled");
        println!("   - Status: PENDING\n");
        println!("8. test_execute_as_user_special_chars: Tests commands with special characters");
        println!("   - Verifies: Proper handling of special characters in commands");
        println!("   - Status: PENDING\n");
        println!("9. test_execute_as_user_invalid_target: Tests execution with invalid target user");
        println!("   - Verifies: Proper error handling for invalid target users");
        println!("   - Status: PENDING\n");
        println!("10. test_execute_as_user_invalid_password: Tests execution with invalid sudo password");
        println!("   - Verifies: Proper error handling for invalid sudo passwords");
        println!("   - Status: PENDING\n");
        println!("11. test_execute_with_custom_id: Tests execute with custom ID");
        println!("   - Verifies: Custom ID is correctly handled");
        println!("   - Status: PENDING\n");
        println!("12. test_execute_as_user_success: Tests successful execution as target user");
        println!("   - Verifies: Successful execution with correct output when running as target user");
        println!("   - Status: PENDING\n");
        println!("Note: All tests are marked as PENDING until they are actually executed.\n");
        println!("To run the tests, execute: cargo test --test ssh_execute_test\n");
        println!("====================================\n");
    }
}